"""
导航数据保存服务
同时保存到数据库和ChromaDB向量数据库
"""

import uuid
import re
from collections import defaultdict
from datetime import datetime
from typing import List, Optional, Dict, Any, Tuple

import jieba
from rank_bm25 import BM25Okapi
from sqlalchemy import and_
from sqlalchemy.orm import Session

from config.chromadb_config import chroma_db
from config.database import get_db_session
from config.logging_config import get_logger
from pojo.ai_web_api_calling import AIWebApiCallingResponse, AIWebApiCallingCreate, AIWebApiCallingUpdate
from pojo.ai_web_api_calling_model import AIWebApiCallingModel

logger = get_logger(__name__)


class NavigationSaveService:
    """导航数据保存服务类"""

    @staticmethod
    def _generate_documents_and_metadata(navigation_data: AIWebApiCallingResponse) -> Tuple[
        List[str], List[Dict[str, Any]], List[str]]:
        """
        生成ChromaDB的文档、元数据和ID
        
        Args:
            navigation_data: 导航数据响应对象
            
        Returns:
            tuple: (documents, metadatas, ids)
        """
        documents = []
        metadatas = []
        ids = []

        # 基础元数据
        base_metadata = {
            "scene_id": navigation_data.scene_id,
            "page_id": navigation_data.page_id,
            "page_name": navigation_data.page_name,
            "navigation_path": navigation_data.navigation_path,
            "tenant_id": navigation_data.tenant_id,
            "created_by": navigation_data.created_by,
            "created_at": navigation_data.created_at.isoformat() if navigation_data.created_at else None,
            "mcp_list_str": ",".join(navigation_data.mcp_list) if navigation_data.mcp_list else ""
        }

        # 处理page_tab
        if navigation_data.page_tab:
            for idx, tab_content in enumerate(navigation_data.page_tab):
                if tab_content.strip():
                    doc_id = f"{navigation_data.scene_id}_{navigation_data.page_id}_tab_{idx}"
                    metadata = {
                        **base_metadata,
                        "content_type": "page_tab",
                        "content_index": idx,
                        "primary_key": f"{navigation_data.scene_id}_{tab_content.strip()}"
                    }
                    documents.append(tab_content.strip())
                    metadatas.append(metadata)
                    ids.append(doc_id)

        # 处理function_desc_list
        if navigation_data.function_desc_list:
            for idx, func_desc in enumerate(navigation_data.function_desc_list):
                if func_desc.strip():
                    doc_id = f"{navigation_data.scene_id}_{navigation_data.page_id}_func_{idx}"
                    metadata = {
                        **base_metadata,
                        "content_type": "function_desc",
                        "content_index": idx,
                        "primary_key": f"{navigation_data.scene_id}_{func_desc.strip()}"
                    }
                    documents.append(func_desc.strip())
                    metadatas.append(metadata)
                    ids.append(doc_id)

        # 处理operation_desc_list
        if navigation_data.operation_desc_list:
            for idx, op_desc in enumerate(navigation_data.operation_desc_list):
                if op_desc.strip():
                    doc_id = f"{navigation_data.scene_id}_{navigation_data.page_id}_op_{idx}"
                    metadata = {
                        **base_metadata,
                        "content_type": "operation_desc",
                        "content_index": idx,
                        "primary_key": f"{navigation_data.scene_id}_{op_desc.strip()}"
                    }
                    documents.append(op_desc.strip())
                    metadatas.append(metadata)
                    ids.append(doc_id)

        return documents, metadatas, ids

    @staticmethod
    def save_navigation_to_database(navigation_data: AIWebApiCallingResponse, db_session: Session) -> bool:
        """
        保存导航数据到数据库
        
        Args:
            navigation_data: 导航数据
            db_session: 数据库会话
            
        Returns:
            bool: 保存是否成功
        """
        try:
            # 创建SQLAlchemy模型实例
            db_item = AIWebApiCallingModel(
                page_id=navigation_data.page_id,
                scene_id=navigation_data.scene_id,
                tenant_id=navigation_data.tenant_id,
                page_name=navigation_data.page_name,
                navigation_path=navigation_data.navigation_path,
                is_deleted=navigation_data.is_deleted,
                created_by=navigation_data.created_by,
                created_at=navigation_data.created_at or datetime.now(),
                updated_by=navigation_data.updated_by,
                updated_at=navigation_data.updated_at
            )

            # 设置JSON字段
            if navigation_data.page_tab:
                db_item.set_page_tab(navigation_data.page_tab)
            if navigation_data.function_desc_list:
                db_item.set_function_desc_list(navigation_data.function_desc_list)
            if navigation_data.operation_desc_list:
                db_item.set_operation_desc_list(navigation_data.operation_desc_list)
            if navigation_data.prev_path_list:
                db_item.set_prev_path_list(navigation_data.prev_path_list)
            if navigation_data.next_path_list:
                db_item.set_next_path_list(navigation_data.next_path_list)
            if hasattr(navigation_data, 'mcp_list') and navigation_data.mcp_list:
                db_item.set_mcp_list(navigation_data.mcp_list)


            # 保存到数据库
            db_session.add(db_item)
            db_session.commit()

            logger.info(f"成功保存导航数据到数据库: {navigation_data.page_id}")
            return True

        except Exception as e:
            logger.error(f"保存导航数据到数据库失败: {e}")
            db_session.rollback()
            return False

    @staticmethod
    def save_navigation_to_chromadb(navigation_data: AIWebApiCallingResponse) -> bool:
        """
        保存导航数据到ChromaDB
        
        Args:
            navigation_data: 导航数据
            
        Returns:
            bool: 保存是否成功
        """
        try:
            if not chroma_db:
                logger.error("ChromaDB 未初始化")
                return False

            # 生成文档、元数据和ID
            documents, metadatas, ids = NavigationSaveService._generate_documents_and_metadata(navigation_data)

            if not documents:
                logger.warning(f"页面 {navigation_data.page_id} 没有可保存的内容到ChromaDB")
                return True

            # 保存到ChromaDB
            success = chroma_db.add_documents(
                documents=documents,
                metadatas=metadatas,
                ids=ids,
                scene_id=navigation_data.scene_id
            )

            if success:
                logger.info(f"成功保存 {len(documents)} 个文档到ChromaDB: {navigation_data.page_id}")
            else:
                logger.error(f"保存文档到ChromaDB失败: {navigation_data.page_id}")

            return success

        except Exception as e:
            logger.error(f"保存导航数据到ChromaDB失败: {e}")
            return False

    @staticmethod
    def update_navigation_in_database(page_id: str, navigation_data: AIWebApiCallingUpdate,
                                      db_session: Session) -> bool:
        """
        更新数据库中的导航数据
        
        Args:
            page_id: 页面ID
            navigation_data: 更新数据
            db_session: 数据库会话
            
        Returns:
            bool: 更新是否成功
        """
        try:
            # 查找现有记录
            db_item = db_session.query(AIWebApiCallingModel).filter(
                AIWebApiCallingModel.page_id == page_id
            ).first()

            if not db_item:
                logger.warning(f"未找到要更新的导航数据: {page_id}")
                return False

            # 更新字段
            if navigation_data.page_name is not None:
                db_item.page_name = navigation_data.page_name
            if navigation_data.navigation_path is not None:
                db_item.navigation_path = navigation_data.navigation_path
            if navigation_data.page_tab is not None:
                db_item.set_page_tab(navigation_data.page_tab)
            if navigation_data.function_desc_list is not None:
                db_item.set_function_desc_list(navigation_data.function_desc_list)
            if navigation_data.operation_desc_list is not None:
                db_item.set_operation_desc_list(navigation_data.operation_desc_list)
            if navigation_data.prev_path_list is not None:
                db_item.set_prev_path_list(navigation_data.prev_path_list)
            if navigation_data.next_path_list is not None:
                db_item.set_next_path_list(navigation_data.next_path_list)
            if hasattr(navigation_data, 'mcp_list') and navigation_data.mcp_list is not None:
                db_item.set_mcp_list(navigation_data.mcp_list)


            db_item.updated_by = navigation_data.updated_by
            db_item.updated_at = datetime.now()

            db_session.commit()

            logger.info(f"成功更新数据库中的导航数据: {page_id}")
            return True

        except Exception as e:
            logger.error(f"更新数据库中的导航数据失败: {e}")
            db_session.rollback()
            return False

    @staticmethod
    def delete_navigation_from_chromadb(page_id: str) -> bool:
        """
        从ChromaDB删除导航数据
        
        Args:
            page_id: 页面ID
            
        Returns:
            bool: 删除是否成功
        """
        try:
            if not chroma_db:
                logger.error("ChromaDB 未初始化")
                return False

            # 删除该页面的所有相关文档，只使用page_id作为过滤条件
            success = chroma_db.delete_documents(
                where={"page_id": page_id}
            )

            if success:
                logger.info(f"成功从ChromaDB删除导航数据: {page_id}")
            else:
                logger.error(f"从ChromaDB删除导航数据失败: {page_id}")

            return success

        except Exception as e:
            logger.error(f"从ChromaDB删除导航数据失败: {e}")
            return False

    @staticmethod
    def delete_navigation_in_database(page_id: str, db_session: Session) -> bool:
        """
        物理删除数据库中的导航数据
        
        Args:
            page_id: 页面ID
            db_session: 数据库会话
            
        Returns:
            bool: 删除是否成功
        """
        try:
            # 查找现有记录
            db_item = db_session.query(AIWebApiCallingModel).filter(
                AIWebApiCallingModel.page_id == page_id
            ).first()

            if not db_item:
                logger.warning(f"未找到要删除的导航数据: {page_id}")
                return False

            # 物理删除
            db_session.delete(db_item)
            db_session.commit()

            logger.info(f"成功物理删除数据库中的导航数据: {page_id}")
            return True

        except Exception as e:
            logger.error(f"物理删除数据库中的导航数据失败: {e}")
            db_session.rollback()
            return False

    @staticmethod
    def _preprocess_text(text: str) -> List[str]:
        """
        文本预处理和分词
        
        Args:
            text: 待处理的文本
            
        Returns:
            List[str]: 分词结果
        """
        if not text:
            return []
        
        # 去除特殊字符，保留中文、英文、数字
        text = re.sub(r'[^\u4e00-\u9fa5a-zA-Z0-9\s]', ' ', text)
        
        # 使用jieba分词
        words = jieba.cut(text.lower().strip())
        
        # 过滤空字符串和单字符
        return [word for word in words if len(word.strip()) > 1]

    @staticmethod
    def _get_all_documents_for_bm25(scene_id: str) -> Tuple[List[Dict[str, Any]], List[List[str]]]:
        """
        获取指定场景下的所有文档用于BM25索引
        
        Args:
            scene_id: 场景ID
            
        Returns:
            Tuple: (文档元数据列表, 分词后的文档列表)
        """
        try:
            if not chroma_db:
                logger.error("ChromaDB 未初始化")
                return [], []

            # 获取该场景下的所有文档
            results = chroma_db.search_by_scene(scene_id=scene_id, limit=1000)
            
            if not results or not isinstance(results, dict):
                logger.error(f"search_by_scene 返回的结果格式不正确: {type(results)}")
                return [], []
            
            if not results.get('documents') or len(results['documents']) == 0:
                return [], []

            documents = results['documents']
            metadatas = results.get('metadatas', [])


            # 确保 documents 和 metadatas 长度一致
            if len(documents) != len(metadatas):
                logger.warning(f"documents 和 metadatas 长度不一致: {len(documents)} vs {len(metadatas)}")
                # 补齐较短的列表
                min_len = min(len(documents), len(metadatas))
                documents = documents[:min_len]
                metadatas = metadatas[:min_len]
                
            # 构建文档信息
            doc_infos = []
            tokenized_docs = []
            
            for i, (doc, metadata) in enumerate(zip(documents, metadatas)):
                # 检查 metadata 是否为字典类型，如果不是则创建默认字典
                if not isinstance(metadata, dict):
                    logger.warning(f"文档 {i} 的元数据不是字典类型: {type(metadata)}, 值: {metadata}")
                    metadata = {}
                
                doc_info = {
                    "content": doc,
                    "metadata": metadata,
                    "page_id": metadata.get("page_id"),
                    "page_name": metadata.get("page_name"),
                    "content_type": metadata.get("content_type"),
                    "primary_key": metadata.get("primary_key"),
                    "navigation_path": metadata.get("navigation_path"),
                }
                doc_infos.append(doc_info)
                
                # 对文档进行分词
                tokens = NavigationSaveService._preprocess_text(doc)
                tokenized_docs.append(tokens)
            
            logger.info(f"获取到 {len(doc_infos)} 个文档用于BM25索引")
            return doc_infos, tokenized_docs
            
        except Exception as e:
            logger.error(f"获取BM25文档失败: {e}")
            return [], []

    @staticmethod
    def search_similar_content(query: str, scene_id: str, limit: int = 10, 
                             vector_weight: float = 0.5, bm25_weight: float = 0.5) -> List[Dict[str, Any]]:
        """
        混合检索：使用向量检索(ChromaDB) + BM25稀疏算法
        
        Args:
            query: 查询内容
            scene_id: 场景ID
            limit: 返回结果数量限制
            vector_weight: 向量检索权重 (默认0.7)
            bm25_weight: BM25检索权重 (默认0.3)
            
        Returns:
            List[Dict]: 混合检索结果
        """
        try:
            if not chroma_db:
                logger.error("ChromaDB 未初始化")
                return []

            # 第一步：向量检索
            vector_results = []
            try:
                chroma_results = chroma_db.query_documents(
                    query_texts=query,
                    scene_id=scene_id,
                    n_results=limit * 2  # 获取更多结果用于混合排序
                )

                if chroma_results.get('documents') and len(chroma_results['documents']) > 0:
                    documents = chroma_results['documents'][0]
                    metadatas = chroma_results.get('metadatas', [[]])[0]
                    distances = chroma_results.get('distances', [[]])[0]

                    for doc, metadata, distance in zip(documents, metadatas, distances):
                        # 检查 metadata 是否为字典类型，如果不是则创建默认字典
                        if not isinstance(metadata, dict):
                            logger.warning(f"向量检索中的元数据不是字典类型: {type(metadata)}, 值: {metadata}")
                            metadata = {}
                        
                        vector_results.append({
                            "content": doc,
                            "metadata": metadata,
                            "vector_score": 1 - distance,  # 距离越小相似度越高
                            "page_id": metadata.get("page_id"),
                            "page_name": metadata.get("page_name"),
                            "content_type": metadata.get("content_type"),
                            "primary_key": metadata.get("primary_key"),
                            "mcp_list": metadata.get("mcp_list_str", "").split(",") if metadata.get("mcp_list_str") else [],
                            "navigation_path": metadata.get("navigation_path"),
                        })

                logger.info(f"向量检索获得 {len(vector_results)} 个结果")
            except Exception as e:
                logger.error(f"向量检索失败: {e}")

            # 第二步：BM25检索
            bm25_results = []
            try:
                # 获取所有文档用于BM25索引
                doc_infos, tokenized_docs = NavigationSaveService._get_all_documents_for_bm25(scene_id)
                
                if tokenized_docs and len(tokenized_docs) > 0:
                    # 创建BM25索引
                    bm25 = BM25Okapi(tokenized_docs)
                    
                    # 对查询进行分词
                    query_tokens = NavigationSaveService._preprocess_text(query)
                    
                    if query_tokens:
                        # 计算BM25分数
                        bm25_scores = bm25.get_scores(query_tokens)
                        
                        # 获取排序后的文档索引
                        sorted_indices = sorted(range(len(bm25_scores)), 
                                              key=lambda i: bm25_scores[i], reverse=True)
                        
                        # 构建BM25结果
                        for idx in sorted_indices[:limit * 2]:  # 获取更多结果用于混合排序
                            if bm25_scores[idx] > 0:  # 只保留有分数的结果
                                doc_info = doc_infos[idx].copy()
                                doc_info["bm25_score"] = float(bm25_scores[idx])
                                bm25_results.append(doc_info)

                logger.info(f"BM25检索获得 {len(bm25_results)} 个结果")
            except Exception as e:
                logger.error(f"BM25检索失败: {e}")

            # 第三步：混合排序
            try:
                # 使用字典合并两种检索结果，以primary_key作为唯一标识
                combined_results = {}
                
                # 添加向量检索结果
                for result in vector_results:
                    key = result.get("primary_key", f"vector_{hash(result['content'])}")
                    combined_results[key] = {
                        **result,
                        "vector_score": result.get("vector_score", 0.0),
                        "bm25_score": 0.0
                    }
                
                # 合并BM25检索结果
                for result in bm25_results:
                    key = result.get("primary_key", f"bm25_{hash(result['content'])}")
                    if key in combined_results:
                        # 如果已存在，更新BM25分数
                        combined_results[key]["bm25_score"] = result.get("bm25_score", 0.0)
                    else:
                        # 如果不存在，添加新结果
                        combined_results[key] = {
                            **result,
                            "vector_score": 0.0,
                            "bm25_score": result.get("bm25_score", 0.0)
                        }

                # 计算混合分数并排序
                final_results = []
                for result in combined_results.values():
                    # 归一化分数（可选）
                    normalized_vector_score = min(result["vector_score"], 1.0)
                    # BM25分数通常需要归一化到0-1范围
                    max_bm25_score = max([r.get("bm25_score", 0) for r in combined_results.values()]) or 1
                    normalized_bm25_score = result["bm25_score"] / max_bm25_score
                    
                    # 计算混合分数
                    hybrid_score = (vector_weight * normalized_vector_score + 
                                  bm25_weight * normalized_bm25_score)
                    
                    result["similarity_score"] = hybrid_score
                    result["vector_score_normalized"] = normalized_vector_score
                    result["bm25_score_normalized"] = normalized_bm25_score
                    
                    final_results.append(result)

                # 按混合分数排序并限制结果数量
                final_results.sort(key=lambda x: x["similarity_score"], reverse=True)
                final_results = final_results[:limit]

                logger.info(f"混合检索最终返回 {len(final_results)} 个结果")
                return final_results

            except Exception as e:
                logger.error(f"混合排序失败: {e}")
                # 如果混合排序失败，返回向量检索结果
                return vector_results[:limit] if vector_results else []

        except Exception as e:
            logger.error(f"混合检索失败: {e}")
            return []

    @staticmethod
    def create_navigation(navigation_data: AIWebApiCallingCreate) -> Optional[AIWebApiCallingResponse]:
        """
        创建导航数据（同时保存到数据库和ChromaDB）
        
        Args:
            navigation_data: 创建数据
            
        Returns:
            Optional[AIWebApiCallingResponse]: 创建的导航数据
        """
        db_session = next(get_db_session())
        try:
            # 生成页面ID
            page_id = str(uuid.uuid4())

            # 构建响应数据
            response_data = AIWebApiCallingResponse(
                page_id=page_id,
                scene_id=navigation_data.scene_id,
                tenant_id=navigation_data.tenant_id,
                page_name=navigation_data.page_name,
                page_tab=navigation_data.page_tab,
                function_desc_list=navigation_data.function_desc_list,
                operation_desc_list=navigation_data.operation_desc_list,
                navigation_path=navigation_data.navigation_path,
                prev_path_list=navigation_data.prev_path_list,
                next_path_list=navigation_data.next_path_list,
                mcp_list=getattr(navigation_data, 'mcp_list', []),
                is_deleted=0,
                created_by=navigation_data.created_by,
                created_at=datetime.now(),
                updated_by=None,
                updated_at=None
            )

            # 保存到数据库
            db_success = NavigationSaveService.save_navigation_to_database(response_data, db_session)
            if not db_success:
                return None

            # 保存到ChromaDB
            chroma_success = NavigationSaveService.save_navigation_to_chromadb(response_data)
            if not chroma_success:
                logger.warning(f"ChromaDB保存失败，但数据库保存成功: {page_id}")

            return response_data

        except Exception as e:
            logger.error(f"创建导航数据失败: {e}")
            return None
        finally:
            db_session.close()

    @staticmethod
    def update_navigation(page_id: str, navigation_data: AIWebApiCallingUpdate) -> Optional[AIWebApiCallingResponse]:
        """
        更新导航数据（同时更新数据库和ChromaDB）
        
        Args:
            page_id: 页面ID
            navigation_data: 更新数据
            
        Returns:
            Optional[AIWebApiCallingResponse]: 更新后的导航数据
        """
        db_session = next(get_db_session())
        try:
            # 更新数据库
            db_success = NavigationSaveService.update_navigation_in_database(page_id, navigation_data, db_session)
            if not db_success:
                return None

            # 获取更新后的数据
            db_item = db_session.query(AIWebApiCallingModel).filter(
                AIWebApiCallingModel.page_id == page_id
            ).first()

            if not db_item:
                return None

            # 转换为响应格式
            response_data = AIWebApiCallingResponse(
                page_id=db_item.page_id,
                scene_id=db_item.scene_id,
                tenant_id=db_item.tenant_id,
                page_name=db_item.page_name,
                page_tab=db_item.get_page_tab(),
                function_desc_list=db_item.get_function_desc_list(),
                operation_desc_list=db_item.get_operation_desc_list(),
                navigation_path=db_item.navigation_path,
                prev_path_list=db_item.get_prev_path_list(),
                next_path_list=db_item.get_next_path_list(),
                mcp_list=db_item.get_mcp_list() if hasattr(db_item, 'get_mcp_list') else [],
                is_deleted=db_item.is_deleted,
                created_by=db_item.created_by,
                created_at=db_item.created_at,
                updated_by=db_item.updated_by,
                updated_at=db_item.updated_at
            )

            # 先删除ChromaDB中的旧数据
            NavigationSaveService.delete_navigation_from_chromadb(page_id)

            # 重新保存到ChromaDB
            chroma_success = NavigationSaveService.save_navigation_to_chromadb(response_data)
            if not chroma_success:
                logger.warning(f"ChromaDB更新失败，但数据库更新成功: {page_id}")

            return response_data

        except Exception as e:
            logger.error(f"更新导航数据失败: {e}")
            return None
        finally:
            db_session.close()

    @staticmethod
    def delete_navigation(page_id: str, updated_by: str) -> bool:
        """
        删除导航数据（物理删除数据库和ChromaDB）
        
        Args:
            page_id: 页面ID
            updated_by: 更新者（保留参数兼容性，但不再使用）
            
        Returns:
            bool: 删除是否成功
        """
        db_session = next(get_db_session())
        try:
            # 获取要删除的数据信息
            db_item = db_session.query(AIWebApiCallingModel).filter(
                AIWebApiCallingModel.page_id == page_id
            ).first()

            if not db_item:
                logger.warning(f"未找到要删除的导航数据: {page_id}")
                return False

            scene_id = db_item.scene_id

            # 物理删除数据库记录
            db_success = NavigationSaveService.delete_navigation_in_database(page_id, db_session)
            if not db_success:
                return False

            # 物理删除ChromaDB记录
            chroma_success = NavigationSaveService.delete_navigation_from_chromadb(page_id)
            if not chroma_success:
                logger.warning(f"ChromaDB删除失败，但数据库物理删除成功: {page_id}")

            return True

        except Exception as e:
            logger.error(f"删除导航数据失败: {e}")
            return False
        finally:
            db_session.close()